import logging
import os
import time

import torch
from _dynamic_dataset.AccCombinationRecord import WriteRecordsToFile



def Algorithm_MR(num_epochs, clients, server, device, combinations, dirPath, acc_record, time_record):
    os.makedirs(dirPath, exist_ok=True)
    logging.basicConfig(filename=f'{dirPath}/main.log', level=logging.INFO)
    # 开始时间
    start_time = time.time()
    for epoch in range(1, num_epochs + 1):
        logging.info(f"epoch={epoch}")
        # 客户端训练
        train_start_time = time.time()
        for client in clients:
            client.train(server.global_model)
        train_end_time = time.time()
        # 这里用(1,)这个组合来代表该轮的贡献度评估时间
        time_record.addCombinationAcc(epoch, (1,), train_end_time - train_start_time)
        # 服务器梯度聚合并计算对应的正确率
        test_start_time = time.time()
        for combination in combinations:
            # 创建客户端子集
            com_clients = []
            for clientId in combination:
                com_clients.append(clients[clientId])
            server.model_aggregate(com_clients, server.sub_model)
            # 在测试集上进行测试
            correct = 0
            total = 0
            server.sub_model.eval()
            with torch.no_grad():
                for data in server.eval_loader:
                    images, labels = data
                    images, labels = images.to(device), labels.to(device)
                    outputs = server.sub_model(images)
                    _, predicted = torch.max(outputs.data, 1)
                    total += labels.size(0)
                    correct += (predicted == labels).sum().item()
            # 真实的正确率
            acc = 100 * correct / total
            acc_record.addCombinationAcc(epoch, combination, acc)
            # print(f"epoch={epoch},com={combination},acc={acc}")
        # 更新全局模型sub_model
        test_end_time = time.time()
        # 这里用(2,)这个组合来代表该轮的贡献度评估时间
        time_record.addCombinationAcc(epoch, (2,), test_end_time - test_start_time)
        logging.info(f"eval_time={test_end_time - test_start_time}")
        server.model_aggregate(clients, server.global_model)
    # 结束时间
    end_time = time.time()
    logging.info(f"总耗时为{end_time - start_time}s")
    WriteRecordsToFile(dirPath + "/acc_records.json", acc_record)
    WriteRecordsToFile(dirPath + "/time_records.json", time_record)

